import pandas as pd
from joblib import dump

from sklearn.ensemble import RandomForestClassifier

# 加载处理过的训练数据
train_data = pd.read_csv('../dataset/merged/train_mini.csv')

# 分离特征和标签
X_train = train_data.drop('Label', axis=1)
y_train = train_data['Label']

# 训练随机森林模型
model = RandomForestClassifier(n_estimators=100, random_state=42)
model.fit(X_train, y_train)

# 保存模型到文件
model_filename = 'detect_model_mini.joblib'
dump(model, model_filename)
print(f"Model saved to {model_filename}")
